import os
import os.path as osp
import shutil
import cv2
import numpy as np
import random
from shapely.geometry import Polygon, MultiPoint  # 多边形


manual_seed = 1
random.seed(manual_seed)
np.random.seed(manual_seed)


BASEDIR = osp.abspath('.')
IMGDIR = osp.join(BASEDIR, 'images')
LABELDIR = osp.join(BASEDIR, 'labelTxt')
OUTDIR = osp.join(BASEDIR, 'neg_patches')
if osp.exists(OUTDIR):
    shutil.rmtree(OUTDIR)
os.makedirs(OUTDIR)


def get_rotated_patch(src_img, ro_rect) -> np.ndarray:
    # cv2.imshow('raw', src_img)
    # cv2.waitKey(1)

    # the order of the box points: bottom left, top left, top right, bottom right
    (x, y), (w, h), a = ro_rect
    rect = ro_rect
    box = cv2.boxPoints(rect)
    box = np.int0(box)

    # print("bounding box: \n{}".format(box))
    # cv2.drawContours(src_img, [box], 0, (255, 0, 255), 3)
    # cv2.imshow('with bbox', src_img)
    # cv2.waitKey(10000)

    # get width and height of the detected rectangle
    width, height = int(w), int(h)

    src_pts = box.astype("float32")
    # coordinate of the points in box points after the rectangle has been straightened
    dst_pts = np.array([[0, height - 1],
                        [0, 0],
                        [width - 1, 0],
                        [width - 1, height - 1]], dtype="float32")

    # the perspective transformation matrix
    M = cv2.getPerspectiveTransform(src_pts, dst_pts)

    # directly warp the rotated rectangle to get the straightened rectangle
    warped = cv2.warpPerspective(src_img, M, (width, height))

    # cv2.imshow('cut result', warped)
    # cv2.waitKey(10000)

    # return cut result
    return warped


def create_a_random_box(box_w, box_h, IMG_W: int, IMG_H: int) -> np.ndarray:  # 返回:四个顶点,旋转矩形
    x = random.randint(0, IMG_W-1)
    y = random.randint(0, IMG_H-1)
    w = box_w
    h = box_h
    angle = random.randint(0, 179)
    ro_rect = ((x, y), (w, h), angle)
    box = cv2.boxPoints(ro_rect)
    xy4 = np.int0(box)
    return xy4, ro_rect


def compute_one_rotation_iou(x_box, y_box):
    # 四边形的二维坐标表示
    xx_box, yy_box = np.array(x_box).reshape(4, 2), np.array(y_box).reshape(4, 2)
    # 构建四边形对象，会自动计算四个点的顺序：左上 左下 右下 右上 左上（返回5个点，最后回到起始点）
    x_poly, y_poly = Polygon(xx_box).convex_hull, Polygon(yy_box).convex_hull

    intersect_area = x_poly.intersection(y_poly).area  # 相交面积
    if intersect_area == 0:
        iou = 0
    else:
        union_area = x_poly.area + y_poly.area - intersect_area  # 总共面积
        iou = intersect_area / union_area
    return iou


def is_overlap(rand_box, positive_areas: list) -> bool: 
    for pos_area in positive_areas:
        if compute_one_rotation_iou(rand_box, pos_area) > 0:
            return True
    return False


key_list = [name.split('.')[-2] for name in os.listdir(IMGDIR) if name.endswith('.png')]
print(key_list)

for k in key_list:
    imgname = k+'.png' 
    labelname = k+'.txt'
    path_img = osp.join(IMGDIR, imgname)
    path_label = osp.join(LABELDIR, labelname)
    
    img = cv2.imread(path_img)
    IMG_H, IMG_W, IMG_C = img.shape
    
    with open(path_label, 'r') as fp:
       lines = fp.readlines()
    # print(lines)
    
    positive_areas = []
    for line in lines:
        one_box_info = line.strip().split()
        *xy4, cls_name, easy_token = one_box_info
        xy4 = list(map(int, xy4))
        xy4 = [[xy4[0],xy4[1]],
               [xy4[2],xy4[3]],
               [xy4[4],xy4[5]],
               [xy4[6],xy4[7]]]
        xy4 = np.array(xy4)  # np.ndarray
        positive_areas.append(xy4)
    
    for i in range(10):  # 每张图里切几块负样本
        rand_xy4, rand_roRect = create_a_random_box(box_w=64, box_h=64, IMG_W=IMG_W, IMG_H=IMG_H)
        times = 0
        while is_overlap(rand_xy4, positive_areas) and times < 10:
            times += 1
            print('collision! try [{}/10] times.'.format(times))
            rand_xy4, rand_roRect = create_a_random_box(box_w=64, box_h=64, IMG_W=IMG_W, IMG_H=IMG_H)
            
        if not is_overlap(rand_xy4, positive_areas):
            patch = get_rotated_patch(img, rand_roRect)
            path_patch_out = osp.join(OUTDIR, f"{k}_bg{i}.png")
            cv2.imwrite(path_patch_out, patch)
            print(f"{k}_bg{i} done.")
            